#!/usr/bin/env python3

import sys
sys.path.append('../../../tools')
from gmcore_plot import *

parser = argparse.ArgumentParser(description='Plot test results')
parser.add_argument('-i', '--input', help='Input file', default='rh.360x180x26.h0.nc')
parser.add_argument('-t', '--time-step', type=int, default=-1, help='Time step to plot')
parser.add_argument('-l', '--loop', action='store_true', help='Loop the plot')
parser.add_argument('-u', '--update-seconds', type=float, default=10, help='Update seconds')
parser.add_argument('-o', '--output', help='Output figure')
args = parser.parse_args()

fig = plt.figure(figsize=(10, 8))
proj = ccrs.PlateCarree(central_longitude=180)

def plot(data_path, fig):
	f = addfile(data_path)

	p = f.ph[args.time_step,:,:,:] * units.Pa

	time0 = parse_datetime(f.time[0])
	time1 = parse_datetime(f.time[args.time_step])
	fig.suptitle(f'Day {(time1 - time0).days}', fontsize=16)

	i = 0

	i += 1; ax = fig.add_subplot(3, 2, i, projection=proj)
	u850 = vinterp_p(p, f.u[args.time_step,:,:,:], 850 * units.hPa)
	plot_contour_map(ax, u850[0,:,:], cmap='jet', left_string='U at 850 hPa', levels=np.linspace(-2, 26, 15))

	i += 1; ax = fig.add_subplot(3, 2, i, projection=proj)
	v850 = vinterp_p(p, f.v[args.time_step,:,:,:], 850 * units.hPa)
	plot_contour_map(ax, v850[0,:,:], cmap='jet', left_string='V at 850 hPa', levels=np.linspace(-16, 16, 17))

	i += 1; ax = fig.add_subplot(3, 2, i, projection=proj)
	plot_contour_map(ax, f.phs[args.time_step,:,:], cmap='jet', left_string='PS', levels=np.linspace(95000, 103000, 17))

	i += 1; ax = fig.add_subplot(3, 2, i, projection=proj)
	t850 = vinterp_p(p, f.t[args.time_step,:,:,:], 850 * units.hPa)
	plot_contour_map(ax, t850[0,:,:], cmap='jet', left_string='T at 850 hPa', levels=np.linspace(281.40, 281.96, 17))

	i += 1; ax = fig.add_subplot(3, 2, i, projection=proj)
	z500 = vinterp_p(p, f.gz[args.time_step,:,:,:], 500 * units.hPa) / 9.80616
	plot_contour_map(ax, z500[0,:,:], cmap='jet', left_string='Z at 500 hPa', levels=np.linspace(5120, 5800, 17))

	i += 1; ax = fig.add_subplot(3, 2, i, projection=proj)
	omg850 = vinterp_p(p, f.omg[args.time_step,:,:,:], 850 * units.hPa)
	plot_contour_map(ax, omg850[0,:,:], cmap='jet', left_string='Omega at 850 hPa')

plot(args.input, fig)
if args.loop:
	loop_plot(args.input, fig, plot, args.update_seconds)
elif args.output:
	plt.savefig(args.output, dpi=300)
else:
	plt.tight_layout()
	plt.show()
